import os
import sys
import queue
import json
from vosk import Model, KaldiRecognizer
import pyaudio
import wave

# 模型路径
MODEL_PATH = os.path.join(os.path.dirname(__file__), 'vosk-model-small-cn-0.22')

# 全局模型单例，避免重复加载
_model_instance = None

def get_model(model_path=MODEL_PATH):
    """
    获取全局唯一的 Vosk Model 实例。
    如果模型尚未加载，则加载一次并缓存。
    后续调用直接返回已加载的模型，避免重复初始化。
    :param model_path: 模型文件夹路径
    :return: Vosk Model 实例
    """
    global _model_instance
    if _model_instance is None:
        if not os.path.exists(model_path):
            print(f"模型目录 {model_path} 不存在，请检查路径！")
            sys.exit(1)
        print("正在加载模型，请稍候...")
        _model_instance = Model(model_path)
        print("模型加载完成！")
    return _model_instance

def recognize_from_microphone(model):
    """
    通过麦克风实时采集音频并进行语音转文字识别。
    适用于需要实时语音识别的场景。
    :param model: 已加载的 Vosk Model 实例（建议通过 get_model 获取）
    """
    SAMPLE_RATE = 16000
    CHUNK = 4000  # 每次读取的帧数
    p = pyaudio.PyAudio()
    stream = p.open(format=pyaudio.paInt16,
                    channels=1,
                    rate=SAMPLE_RATE,
                    input=True,
                    frames_per_buffer=CHUNK)
    stream.start_stream()
    rec = KaldiRecognizer(model, SAMPLE_RATE)
    print("请开始说话，按 Ctrl+C 退出...")
    try:
        while True:
            data = stream.read(CHUNK, exception_on_overflow=False)
            if rec.AcceptWaveform(data):
                result = rec.Result()
                text = json.loads(result).get('text', '')
                if text.strip():
                    print(f"识别结果: {text}")
            else:
                partial = rec.PartialResult()
                partial_text = json.loads(partial).get('partial', '')
                if partial_text.strip():
                    print(f"中间结果: {partial_text}", end='\r')
    except KeyboardInterrupt:
        print("\n识别结束。")
    finally:
        stream.stop_stream()
        stream.close()
        p.terminate()

def recognize_from_file(model, wav_path):
    """
    读取指定的音频文件（需为16kHz、16bit、单声道wav格式），并进行语音转文字识别。
    适用于批量处理或离线音频文件识别的场景。
    :param model: 已加载的 Vosk Model 实例（建议通过 get_model 获取）
    :param wav_path: 音频文件路径
    """
    if not os.path.exists(wav_path):
        print(f"音频文件 {wav_path} 不存在！")
        return
    wf = wave.open(wav_path, "rb")
    if wf.getnchannels() != 1 or wf.getsampwidth() != 2 or wf.getframerate() != 16000:
        print("请提供 16kHz, 16bit, 单声道 wav 文件！")
        wf.close()
        return
    rec = KaldiRecognizer(model, wf.getframerate())
    print("正在识别音频文件...")
    while True:
        data = wf.readframes(4000)
        if len(data) == 0:
            break
        if rec.AcceptWaveform(data):
            result = rec.Result()
            text = json.loads(result).get('text', '')
            if text.strip():
                print(f"识别结果: {text}")
        else:
            partial = rec.PartialResult()
            partial_text = json.loads(partial).get('partial', '')
            if partial_text.strip():
                print(f"中间结果: {partial_text}", end='\r')
    # 打印最终结果
    final = rec.FinalResult()
    text = json.loads(final).get('text', '')
    if text.strip():
        print(f"最终识别结果: {text}")
    wf.close()

def recognize_wav_file_return_text(model, wav_path):
    """
    读取指定的音频文件（需为16kHz、16bit、单声道wav格式），并返回语音转文字的最终文本。
    :param model: 已加载的 Vosk Model 实例（建议通过 get_model 获取）
    :param wav_path: 音频文件路径
    :return: 识别到的文本字符串
    """
    if not os.path.exists(wav_path):
        return None
    wf = wave.open(wav_path, "rb")
    if wf.getnchannels() != 1 or wf.getsampwidth() != 2 or wf.getframerate() != 16000:
        wf.close()
        return None
    rec = KaldiRecognizer(model, wf.getframerate())
    while True:
        data = wf.readframes(4000)
        if len(data) == 0:
            break
        rec.AcceptWaveform(data)
    final = rec.FinalResult()
    text = json.loads(final).get('text', '')
    wf.close()
    return text

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Vosk 语音识别示例")
    parser.add_argument('--mode', choices=['mic', 'file'], help='mic: 读取麦克风; file: 读取音频文件')
    parser.add_argument('--wav', type=str, help='音频文件路径，仅在 file 模式下需要')
    parser.add_argument('--model', type=str, default=MODEL_PATH, help='模型文件夹路径，可选，默认使用 vosk-model-small-cn-0.22')
    args = parser.parse_args()

    # 如果未指定 mode，则与用户交互选择
    mode = args.mode
    wav_path = args.wav
    if not mode:
        print("请选择启动模式：")
        print("1. 读取麦克风 (mic)")
        print("2. 读取音频文件 (file)")
        while True:
            choice = input("请输入数字 1 或 2 并回车: ").strip()
            if choice == '1':
                mode = 'mic'
                break
            elif choice == '2':
                mode = 'file'
                break
            else:
                print("输入无效，请重新输入 1 或 2。")
    if mode == 'file' and not wav_path:
        wav_path = input("请输入音频文件路径: ").strip()

    model = get_model(args.model)

    if mode == 'mic':
        recognize_from_microphone(model)
    elif mode == 'file':
        if not wav_path:
            print("未指定音频文件路径，程序退出。")
            sys.exit(1)
        recognize_from_file(model, wav_path)

